#!/usr/bin/python3
# -*- coding: utf-8 -*-

#publish images
import rospy
from sensor_msgs.msg import Image
import gxipy as gx
from cv_bridge import CvBridge,CvBridgeError
import cv2
import sys
import traceback
import time
import os
import numpy as np
import csv

np.set_printoptions(formatter={'float': '{: 0.4f}'.format})

import message_filters

from geometry_msgs.msg import Transform,Pose,Point32

from bs_imu import bs_imu
from bs_img import bs_img_real
from bs_img import bs_img_base
from bs_img.bs_cfg_real_5m import *
from bs_img.bs_lie_solve import *
from bs_img.serial_pack import *


class BS_Solver:
    def __init__(self):
        self.invalid_frame_num_drogue = 0
        self.invalid_frame_max = 8     #if invaild frame < invalid_frame_max, output results of the last frame

        #Kalman Filter
        self.Q = np.eye(3)
        self.R = np.eye(3)
        self.x_drogue = np.array([np.nan, np.nan, np.nan])
        self.P_drogue = np.eye(3)


        self.wa_drogue_pub = rospy.Publisher("/bs_debuger/result/wa_drogue", Transform, queue_size=1)
        
        self.wa_img_pub = rospy.Publisher('image_dh/wa', Image, queue_size=1)
        self.lf_drogue_pub = rospy.Publisher("/bs_debuger/result/wa_drogue", Transform, queue_size=1)
        
        self.bridge = CvBridge()

        
        device_manager = gx.DeviceManager()
        dev_num, dev_info_list = device_manager.update_device_list()
        if dev_num == 0:
            sys.exit(1)
        '''
        每个相机的sn都是独特的，在相机下面写着，这样就可以利用这个分开长短焦相机
        sn作为launch文件中外传入的参数
        '''

        param_list = rospy.get_param_names()
        for param in param_list:  #反证只有一个键，虽说跟node有关
            if '/sn' in param:
                sn = rospy.get_param(param) 
            if '/ExposureTime' in param:
                ExposureTime = rospy.get_param(param)
        self.cam = device_manager.open_device_by_sn(sn)
        #set parameters
        self.cam.TriggerMode.set(gx.GxSwitchEntry.OFF)#set continuous acquisition
        self.cam.ExposureTime.set(ExposureTime)  #set exposure time
        self.cam.stream_on()

        csvfile = open ("/home/nuc001/bs_gd_single-master/bag/data_wa.csv","w") 
        self.writer = csv.writer(csvfile)
        self.writer.writerow(['cam_cls','target_cls','IsValid', 'time','x','y','z'])
        
        self.solve()
        
        if rospy.is_shutdown():
            self.cam.stream_off()
            self.cam.close_device()

    def wa_img_cb(self,img_gray):
        _, img_bin = cv2.threshold(img_gray, 50, 0xff, cv2.THRESH_BINARY)
        # 如果相机倒置需要旋转一下
        # lf_img_gray = cv2.rotate(lf_img_gray, cv2.ROTATE_180)

        drogue_t_vec = [np.nan]*3
        
        gt_pts,_ = bs_img_real.img_to_pts(img_bin)
        
        len_gt_pts = len(gt_pts)
        if len_gt_pts >= 18:    #噪点太多（可能逆光），为了防止其解算时间太长耽误后面的进度，这里直接跳过
            return
        
        # 求解锥套位置
        drogue_result_dict = bs_img_real.solve_drogue_wa(gt_pts, mapping_coeffs, distortion_center, stretch_matrix, taper_radius=r_drogue)
            
        drogue_t_vec = drogue_result_dict["t_vec"]
        
        drogue_msg = Transform()
        drogue_msg.translation.x = drogue_t_vec[0]
        drogue_msg.translation.y = drogue_t_vec[1]
        drogue_msg.translation.z = drogue_t_vec[2]
        self.wa_drogue_pub.publish(drogue_msg)

        drogue_is_visiable = True if not(np.all(np.isnan(drogue_t_vec))) else False

        # # KF for drogue
        # if not(drogue_is_visiable):
        #     self.invalid_frame_num_drogue = self.invalid_frame_num_drogue + 1
        #     if self.invalid_frame_num_drogue <= self.invalid_frame_max:
        #         self.drogue_res = self.x_drogue
        #         if not(np.all(np.isnan(self.drogue_res))):
        #             drogue_res_frame = serial_frame("drogue", True, round((self.drogue_res[2])*100),round(self.drogue_res[0]*100),round((self.drogue_res[1])*100))  
        #         else:
        #             drogue_res_frame = serial_frame("drogue",False,round(0),round(0),round(0))
        #     else:
        #         self.x_drogue = np.array([np.nan, np.nan, np.nan])
        #         drogue_res_frame = serial_frame("drogue",False,round(0),round(0),round(0))
        
        # else:
        #     self.invalid_frame_num_drogue = 0

        #     if np.all(np.isnan(self.x_drogue)):   #the first frame
        #         self.x_drogue = drogue_t_vec
        #     else:
        #         z_measure_drogue = np.array(drogue_t_vec)
        #         self.P_drogue = self.P_drogue + self.Q
        #         K = np.matmul(self.P_drogue, np.linalg.pinv(self.P_drogue + self.R))
        #         self.x_drogue = self.x_drogue + np.matmul(K, z_measure_drogue - self.x_drogue)
        #         self.P_drogue = np.matmul(np.eye(3) - K, self.P_drogue)
        #     drogue_res_frame = serial_frame("drogue",True,round((self.x_drogue[2])*100),round(self.x_drogue[0]*100),round((self.x_drogue[1])*100))  

        # uart.write(drogue_res_frame)
        # time.sleep(0.0005)

        # print(drogue_t_vec)
        if drogue_is_visiable:
            cv2.putText(img_gray, "wa_drogue:[{:.2f},{:.2f},{:.2f}]".format(drogue_t_vec[2],drogue_t_vec[0],drogue_t_vec[1]), (100,300),cv2.FONT_HERSHEY_SIMPLEX, 4, 255, thickness=5)
            self.writer.writerow(["wa","drogue",True, time.time(), drogue_t_vec[0], drogue_t_vec[1], drogue_t_vec[2]])
        else:
            cv2.putText(img_gray, "wa_drogue:Invalid", (100,300),cv2.FONT_HERSHEY_SIMPLEX, 4, 255, thickness=5)
            self.writer.writerow(["wa","drogue",False, time.time(), 0, 0, 0])
        
        # cv2.imwrite(f"/home/nuc001/bs_gd_single-master/bag/image_wa/{time.time()}.jpg",img_gray)
        msg=self.bridge.cv2_to_imgmsg(img_gray,"mono8")   #bgr image
        self.wa_img_pub.publish(msg)


    def solve(self):
        while not rospy.is_shutdown():
            try:
                if not os.path.exists("/home/nuc001/bs_gd_single-master/bag/image_wa"):
                    os.makedirs("/home/nuc001/bs_gd_single-master/bag/image_wa")
                raw_image = self.cam.data_stream[0].get_image()
                numpy_image = raw_image.get_numpy_array()
                if numpy_image is None:
                    continue


                # if not os.path.exists("/home/brightsoul/bs_gd_single/bag/image_wa"):
                #     os.makedirs("/home/brightsoul/bs_gd_single/bag/image_wa")          
                # cv2.imwrite(f"/home/brightsoul/bs_gd_single/bag/image_wa/{time.time()}.jpg",numpy_image)
                
                self.wa_img_cb(numpy_image)

                msg=self.bridge.cv2_to_imgmsg(numpy_image,"mono8")   #bgr image
                self.wa_img_pub.publish(msg)
            except Exception as e:
                print(traceback.format_exc())
            
            
       

if __name__ == '__main__':
    rospy.init_node('bs_solver_node', anonymous=True)

    # port, baud =  '/dev/ttyUSB0', 115200
    # uart = serial.Serial(port,baud,parity=serial.PARITY_NONE)
    bs_solver = BS_Solver()

    rospy.spin()